{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Overview\n",
    "\n",
    "This notebook can be used with both a single or multi- speaker corpus and allows the interactive plotting of speaker embeddings linked to underlying audio (see instructions in the repo's speaker_embedding directory)\n",
    "\n",
    "Depending on the directory structure used for your corpus, you may need to adjust handling of **speaker_to_utter** and **locations**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import random\n",
    "from collections import defaultdict\n",
    "from pathlib import Path\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import umap\n",
    "\n",
    "from TTS.bin.compute_embeddings import compute_embeddings\n",
    "from TTS.config import load_config\n",
    "from TTS.config.shared_configs import BaseDatasetConfig\n",
    "from TTS.tts.datasets import load_tts_samples\n",
    "from TTS.utils.audio import AudioProcessor\n",
    "\n",
    "from bokeh.io import output_notebook, show\n",
    "from bokeh.plotting import figure\n",
    "from bokeh.models import HoverTool, ColumnDataSource, BoxZoomTool, ResetTool, OpenURL, TapTool\n",
    "from bokeh.transform import factor_cmap\n",
    "from bokeh.palettes import Category10\n",
    "\n",
    "output_notebook()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For larger sets of speakers, you can use `Category20`, but you need to change it in the `pal` variable too\n",
    "\n",
    "List of Bokeh palettes here: https://docs.bokeh.org/en/latest/docs/reference/palettes.html\n",
    "\n",
    "**NB:** if you have problems with other palettes, first see https://stackoverflow.com/questions/48333820/why-do-some-bokeh-palettes-raise-a-valueerror-when-used-in-factor-cmap"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Config\n",
    "\n",
    "You should adjust all the paths to point at the relevant locations for you locally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dataset\n",
    "formatter_name = \"ljspeech\"\n",
    "dataset_name = \"ljspeech\"\n",
    "dataset_path = \"path/to/LJSpeech-1.1\"\n",
    "meta_file_train = \"metadata.csv\"\n",
    "\n",
    "# Speaker encoder\n",
    "se_model_path = \"https://github.com/coqui-ai/TTS/releases/download/speaker_encoder_model/model_se.pth.tar\"\n",
    "se_config_path = \"https://github.com/coqui-ai/TTS/releases/download/speaker_encoder_model/config_se.json\"\n",
    "embedding_path = \"speakers.pth\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_config = BaseDatasetConfig()\n",
    "dataset_config.formatter = formatter_name\n",
    "dataset_config.dataset_name = dataset_name\n",
    "dataset_config.path = dataset_path\n",
    "dataset_config.meta_file_train = meta_file_train\n",
    "\n",
    "meta_data_train, meta_data_eval = load_tts_samples(dataset_config, eval_split=False)\n",
    "utt_to_wav = {\n",
    "    item[\"audio_unique_name\"]: str(Path(item[\"audio_file\"]).relative_to(dataset_path)) for item in meta_data_train\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute embeddings\n",
    "\n",
    "You can skip this if you have already computed embeddings with `TTS/bin/compute_embeddings.py`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "compute_embeddings(\n",
    "    model_path=se_model_path,\n",
    "    config_path=se_config_path,\n",
    "    output_path=embedding_path,\n",
    "    formatter_name=formatter_name,\n",
    "    dataset_name=dataset_name,\n",
    "    dataset_path=dataset_path,\n",
    "    meta_file_train=meta_file_train,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot Umap"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Bring in the embeddings created by `TTS/bin/compute_embeddings.py`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings = torch.load(embedding_path, weights_only=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "speakers = set()\n",
    "speaker_to_utter = defaultdict(list)\n",
    "for idx, embedding in embeddings.items():\n",
    "    speaker = embedding[\"name\"]\n",
    "    speakers.add(speaker)\n",
    "    speaker_to_utter[speaker].append(idx)\n",
    "print(f\"Speaker count: {len(speakers)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up the embeddings\n",
    "\n",
    "Adjust the number of speakers to select and the number of utterances from each speaker and they will be randomly sampled from the corpus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeds = []\n",
    "labels = []\n",
    "locations = []\n",
    "\n",
    "# single speaker\n",
    "num_speakers = 1\n",
    "num_utters = 1000\n",
    "\n",
    "# multi speaker\n",
    "# num_speakers = 10\n",
    "# num_utters = 20\n",
    "\n",
    "speaker_idxs = random.sample(list(speakers), num_speakers)\n",
    "\n",
    "for speaker_num, speaker_idx in enumerate(speaker_idxs):\n",
    "    speakers_utter = speaker_to_utter[speaker_idx]\n",
    "    utter_idxs = random.sample(speakers_utter, num_utters)\n",
    "    for utter_idx in utter_idxs:\n",
    "        embed = np.array(embeddings[utter_idx][\"embedding\"])\n",
    "        embeds.append(embed)\n",
    "        labels.append(speaker_idx)\n",
    "        locations.append(utt_to_wav[utter_idx])\n",
    "embeds = np.stack(embeds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load embeddings with UMAP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = umap.UMAP()\n",
    "projection = model.fit_transform(embeds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interactively charting the data in Bokeh\n",
    "\n",
    "You can use the regular Bokeh [tools](https://docs.bokeh.org/en/latest/docs/user_guide/interaction/tools.html) to explore the data, with reset setting it back to normal\n",
    "\n",
    "Once you have started the local server (see cell below) you can then click on plotted points which will open a tab to play the audio for that point, enabling easy exploration of your corpus\n",
    "\n",
    "File location in the tooltip is given relative to **AUDIO_PATH**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "source_wav_stems = ColumnDataSource(\n",
    "    data=dict(\n",
    "        x=projection.T[0].tolist(),\n",
    "        y=projection.T[1].tolist(),\n",
    "        desc=locations,\n",
    "        label=labels,\n",
    "    )\n",
    ")\n",
    "\n",
    "hover = HoverTool(tooltips=[(\"file\", \"@desc\"), (\"speaker\", \"@label\")])\n",
    "\n",
    "### Optionally consider adding these to the tooltips if you want additional detail\n",
    "# for the coordinates: (\"(x,y)\", \"($x, $y)\"),\n",
    "# for the index of the embedding / wav file: (\"index\", \"$index\"),\n",
    "\n",
    "factors = list(set(labels))\n",
    "pal_size = max(len(factors), 3)\n",
    "pal = Category10[pal_size]\n",
    "\n",
    "p = figure(width=600, height=400, tools=[hover, BoxZoomTool(), ResetTool(), TapTool()])\n",
    "p.scatter(\n",
    "    \"x\",\n",
    "    \"y\",\n",
    "    source=source_wav_stems,\n",
    "    color=factor_cmap(\"label\", palette=pal, factors=factors),\n",
    ")\n",
    "\n",
    "url = \"http://localhost:8000/@desc\"\n",
    "taptool = p.select(type=TapTool)\n",
    "taptool.callback = OpenURL(url=url)\n",
    "\n",
    "show(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Local server to serve wav files from corpus\n",
    "\n",
    "This is required so that when you click on a data point the hyperlink associated with it will be served the file locally.\n",
    "\n",
    "There are other ways to serve this if you prefer and you can also run the commands manually on the command line\n",
    "\n",
    "The server will continue to run until stopped. To stop it simply interupt the kernel (ie square button or under Kernel menu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd $dataset_path\n",
    "%pwd\n",
    "!python -m http.server"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
